from pymilvus import connections, Collection, FieldSchema, CollectionSchema, DataType

# 连接到Milvus服务器
connections.connect(
    uri="http://111.229.125.24:19530",
    db_name="customer_service",
    token="root:Milvus"
)

# 定义collection的字段
general_qa_fields = [
    FieldSchema(name="faq_id", dtype=DataType.INT64, is_primary=True),
    
    # 向量字段 - 支持多种问法
    FieldSchema(name="question_embedding", dtype=DataType.FLOAT_VECTOR, dim=384),
    FieldSchema(name="keywords_embedding", dtype=DataType.FLOAT_VECTOR, dim=384),
    
    # 内容字段
    FieldSchema(name="question", dtype=DataType.VARCHAR, max_length=500),
    FieldSchema(name="answer", dtype=DataType.VARCHAR, max_length=2000),
    FieldSchema(name="keywords", dtype=DataType.ARRAY, element_type=DataType.VARCHAR, max_capacity=20, max_length=100),
    FieldSchema(name="category", dtype=DataType.VARCHAR, max_length=50),
    
    # 场景标签
    FieldSchema(name="scene_tags", dtype=DataType.ARRAY, element_type=DataType.VARCHAR, max_capacity=10, max_length=50),
    FieldSchema(name="emotion_type", dtype=DataType.VARCHAR, max_length=20),
    
    # 优先级和统计
    FieldSchema(name="priority", dtype=DataType.INT8),
    FieldSchema(name="use_count", dtype=DataType.INT32),
    FieldSchema(name="satisfaction_rate", dtype=DataType.FLOAT),
    FieldSchema(name="update_time", dtype=DataType.INT64)
]

# 创建collection schema
general_qa_schema = CollectionSchema(
    fields=general_qa_fields,
    description="General customer service Q&A and templates"
)

# 创建collection
collection_name = "general_faq"
collection = Collection(
    name=collection_name,
    schema=general_qa_schema,
    using='default'
)

# 为向量字段创建索引
index_params = {
    "metric_type": "L2",
    "index_type": "IVF_FLAT",
    "params": {"nlist": 1024}
}

collection.create_index(field_name="question_embedding", index_params=index_params)
collection.create_index(field_name="keywords_embedding", index_params=index_params)

# 加载collection到内存
collection.load()

print(f"Collection {collection_name} created successfully!")

# 关闭连接
connections.disconnect("default")
